from get_index import make_oracle, get_index
import numpy as np
import random
from utils import load_model

def model_predict(model,input_data, start_point):
    n = len(input_data)
    xy = np.array(input_data)

    # 归一化处理
    xy_min = xy.min(axis=0)
    xy_max = xy.max(axis=0)
    xy = (xy - xy_min) / (xy_max - xy_min)

    oracle = make_oracle(model, xy, temperature=1.0)
    
    start_index = input_data.index(start_point)
    
    tour = get_index(oracle, n)
    tour = list(map(int, tour))

    start_pos = tour.index(start_index)
    tour = tour[start_pos:] + tour[:start_pos]
    sorted_data = [input_data[i] for i in tour]
    return sorted_data


def generate_random_input_test(n_points=1000):
    """生成一个包含 n_points 个随机点的列表，坐标范围 [0, 1]"""
    return [[random.uniform(0, 1), random.uniform(0, 1)] for _ in range(n_points)]

if __name__ == '__main__':
    # 生成 100 个测试用例
    model = load_model("params/tsp_100",device=0)[0]
    test_cases = [generate_random_input_test(n_points=20) for _ in range(100)]
    
    for idx, input_data in enumerate(test_cases):
        start_point = random.choice(input_data)  # 随机选择起点
        print(f"Running test case {idx + 1}")
        tour = model_predict(model,input_data, start_point)
        print(f"Tour for test case {idx + 1}: {tour[:10]}...")  # 只打印前10个点以简洁显示

